use crate::compile::benchmark::codegen_backend::CodegenBackend;
use crate::compile::benchmark::profile::Profile;
use anyhow::{anyhow, Context};
use log::debug;
use reqwest::StatusCode;
use std::ffi::OsStr;
use std::fs;
use std::io::{BufReader, Read};
use std::path::{Path, PathBuf};
use std::process::Command;
use std::{fmt, str};
use tar::Archive;
use xz2::bufread::XzDecoder;

pub enum SysrootDownloadError {
    SysrootShaNotFound,
    IO(anyhow::Error),
}

impl SysrootDownloadError {
    pub fn as_anyhow_error(self) -> anyhow::Error {
        match self {
            SysrootDownloadError::SysrootShaNotFound => {
                anyhow::anyhow!("Sysroot was not found on CI")
            }
            SysrootDownloadError::IO(error) => error,
        }
    }
}

/// Sysroot downloaded from CI.
pub struct Sysroot {
    sha: String,
    pub components: ToolchainComponents,
    triple: String,
    preserve: bool,
}

impl Sysroot {
    pub async fn install(
        cache_directory: &Path,
        sha: String,
        triple: &str,
        backends: &[CodegenBackend],
    ) -> Result<Self, SysrootDownloadError> {
        // The structure of this directory is load-bearing.
        // We use the commit SHA as the top-level key, to have a quick way of estimating how many
        // toolchains have been installed in the cache directory.
        // We also use a nested directory below the target tuple, because rustc outputs weird things
        // when we query it with `--print sysroot` and its sysroot is located in a directory that
        // corresponds to a valid target name.
        let cache_directory = cache_directory.join(&sha).join(triple).join("toolchain");
        fs::create_dir_all(&cache_directory).map_err(|e| SysrootDownloadError::IO(e.into()))?;

        let download = SysrootDownload {
            cache_directory: cache_directory.clone(),
            rust_sha: sha.clone(),
            triple: triple.to_owned(),
        };

        let requires_cranelift = backends.contains(&CodegenBackend::Cranelift);

        let stamp = SysrootStamp::load_from_dir(&cache_directory);
        match stamp {
            Ok(stamp) => {
                log::debug!("Found existing stamp for {sha}/{triple}: {stamp:?}");
                // We should already have a complete sysroot present on disk, check if we need to
                // download optional components
                if requires_cranelift && !stamp.cranelift {
                    download.get_and_extract(Component::Cranelift).await?;
                }
            }
            Err(_) => {
                log::debug!(
                    "No existing stamp found for {sha}/{triple}, downloading a fresh sysroot"
                );

                // There is no stamp, download everything
                download.get_and_extract(Component::Rustc).await?;
                download.get_and_extract(Component::Std).await?;
                download.get_and_extract(Component::Cargo).await?;
                download.get_and_extract(Component::RustSrc).await?;
                if requires_cranelift {
                    download.get_and_extract(Component::Cranelift).await?;
                }
            }
        }

        // Update the stamp
        let stamp = SysrootStamp {
            cranelift: requires_cranelift,
        };
        stamp
            .store_to_dir(&cache_directory)
            .map_err(SysrootDownloadError::IO)?;

        let sysroot = download.into_sysroot().map_err(SysrootDownloadError::IO)?;

        Ok(sysroot)
    }

    pub fn preserve(&mut self) {
        self.preserve = true;
    }
}

impl Drop for Sysroot {
    fn drop(&mut self) {
        if self.preserve {
            return;
        }
        fs::remove_dir_all(format!("cache/{}", self.sha)).unwrap_or_else(|err| {
            log::info!(
                "failed to remove {:?}, please do so manually: {:?}",
                format!("cache/{}", self.sha),
                err
            );
        });
    }
}

const SYSROOT_STAMP_FILENAME: &str = ".sysroot-stamp.json";

/// Stores a proof on disk that we have downloaded a complete sysroot.
/// It is used to avoid redownloading a sysroot if it already exists on disk.
#[derive(serde::Serialize, serde::Deserialize, Debug)]
struct SysrootStamp {
    /// Was Cranelift downloaded as a part of the sysroot?
    cranelift: bool,
}

impl SysrootStamp {
    fn load_from_dir(dir: &Path) -> anyhow::Result<Self> {
        let data = std::fs::read(dir.join(SYSROOT_STAMP_FILENAME))?;
        let stamp: SysrootStamp = serde_json::from_slice(&data)?;
        Ok(stamp)
    }

    fn store_to_dir(&self, dir: &Path) -> anyhow::Result<()> {
        let file = std::fs::File::create(dir.join(SYSROOT_STAMP_FILENAME))?;
        Ok(serde_json::to_writer(file, self)?)
    }
}

#[derive(Debug, Clone)]
struct SysrootDownload {
    cache_directory: PathBuf,
    rust_sha: String,
    triple: String,
}

const BASE_URL: &str = "https://ci-artifacts.rust-lang.org/rustc-builds";

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
enum Component {
    Cargo,
    Rustc,
    Std,
    RustSrc,
    Cranelift,
}

impl fmt::Display for Component {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            Component::Cargo => write!(f, "cargo"),
            Component::Rustc => write!(f, "rustc"),
            Component::Std => write!(f, "rust-std"),
            Component::RustSrc => write!(f, "rust-src"),
            Component::Cranelift => write!(f, "rustc-codegen-cranelift"),
        }
    }
}

impl Component {
    fn url(&self, channel: &str, sysroot: &SysrootDownload, triple: &str) -> String {
        let suffix = if *self == Component::RustSrc {
            String::new()
        } else {
            format!("-{triple}")
        };
        format!(
            "{base}/{sha}/{module}-{channel}{suffix}.tar.xz",
            base = BASE_URL,
            module = self,
            sha = sysroot.rust_sha,
            channel = channel,
            suffix = suffix,
        )
    }
}

impl SysrootDownload {
    fn into_sysroot(self) -> anyhow::Result<Sysroot> {
        let sysroot_bin_dir = self.cache_directory.join("bin");
        let sysroot_bin = |name| {
            let path = sysroot_bin_dir.join(name);
            path.canonicalize().with_context(|| {
                format!(
                    "failed to canonicalize {} path for {}: {:?}",
                    name, self.rust_sha, path
                )
            })
        };

        let host_libdir = self.cache_directory.join("lib");
        let target_libdir = target_libdir_from_host_libdir(&host_libdir, &self.triple);
        let components = ToolchainComponents::from_binaries_and_libdir(
            sysroot_bin("rustc")?,
            Some(sysroot_bin("rustdoc")?),
            sysroot_bin("clippy-driver").ok(),
            sysroot_bin("cargo")?,
            &host_libdir,
            &target_libdir,
        )?;

        Ok(Sysroot {
            components,
            sha: self.rust_sha,
            triple: self.triple,
            preserve: false,
        })
    }

    async fn get_and_extract(&self, component: Component) -> Result<(), SysrootDownloadError> {
        // We usually have nightlies but we want to avoid breaking down if we
        // accidentally end up with a beta or stable commit.
        let urls = [
            component.url("nightly", self, &self.triple),
            component.url("beta", self, &self.triple),
            component.url("stable", self, &self.triple),
        ];

        // Did we see any other error than 404?
        let mut found_error_that_is_not_404 = false;
        for url in &urls {
            log::debug!("requesting: {}", url);
            let resp = reqwest::get(url)
                .await
                .map_err(|e| SysrootDownloadError::IO(e.into()))?;
            log::debug!("response status: {}", resp.status());

            match resp.status() {
                s if s.is_success() => {
                    let bytes: Vec<u8> = resp
                        .bytes()
                        .await
                        .map_err(|e| SysrootDownloadError::IO(e.into()))?
                        .into();
                    let reader = XzDecoder::new(BufReader::new(bytes.as_slice()));
                    match self.extract(component, reader) {
                        Ok(()) => return Ok(()),
                        Err(err) => {
                            log::warn!("extracting {url} failed: {err:?}");
                            found_error_that_is_not_404 = true;
                        }
                    }
                }
                StatusCode::NOT_FOUND => {}
                _ => {
                    log::error!("response body: {}", resp.text().await.unwrap_or_default());
                    found_error_that_is_not_404 = true
                }
            }
        }

        if !found_error_that_is_not_404 {
            // The only errors we saw were 404, so we assume that the toolchain is simply not on CI
            Err(SysrootDownloadError::SysrootShaNotFound)
        } else {
            Err(SysrootDownloadError::IO(anyhow!(
                "unable to download sha {} triple {} module {component} from any of {urls:?}",
                self.rust_sha,
                self.triple,
            )))
        }
    }

    fn extract<T: Read>(&self, component: Component, reader: T) -> anyhow::Result<()> {
        let mut archive = Archive::new(reader);
        let prefix = match component {
            Component::Std => format!("rust-std-{}", self.triple),
            Component::Cranelift => format!("{component}-preview"),
            _ => component.to_string(),
        };

        let unpack_into = &self.cache_directory;

        for entry in archive.entries()? {
            let mut entry = entry?;
            let path = entry.path()?.into_owned();
            let mut components = path.components();
            assert!(components.next().is_some(), "strip container directory");
            let path = components.as_path();

            let path = if let Ok(path) = path.strip_prefix(&prefix) {
                unpack_into.join(path)
            } else {
                continue;
            };
            fs::create_dir_all(path.parent().unwrap()).with_context(|| {
                format!(
                    "could not create intermediate directories for {}",
                    path.display()
                )
            })?;
            entry.unpack(path)?;
        }

        Ok(())
    }
}

fn target_libdir_from_host_libdir(dir: &Path, target: &str) -> PathBuf {
    dir.join("rustlib").join(target).join("lib")
}

/// Representation of a toolchain that can be used to compile Rust programs.
#[derive(Debug, Clone)]
pub struct Toolchain {
    pub components: ToolchainComponents,
    pub id: String,
    pub triple: String,
}

impl Toolchain {
    pub fn from_sysroot(sysroot: &Sysroot, id: String) -> Self {
        Self {
            components: sysroot.components.clone(),
            id,
            triple: sysroot.triple.clone(),
        }
    }
}

#[derive(Debug, Clone, Default)]
pub struct ToolchainComponents {
    pub rustc: PathBuf,
    pub rustdoc: Option<PathBuf>,
    pub clippy: Option<PathBuf>,
    pub cargo: PathBuf,
    pub cargo_configs: Vec<String>,
    pub lib_rustc: Option<PathBuf>,
    pub lib_std: Option<PathBuf>,
    pub lib_llvm: Option<PathBuf>,
}

impl ToolchainComponents {
    fn from_binaries_and_libdir(
        rustc: PathBuf,
        rustdoc: Option<PathBuf>,
        clippy: Option<PathBuf>,
        cargo: PathBuf,
        host_libdir: &Path,
        target_libdir: &Path,
    ) -> anyhow::Result<Self> {
        let mut component = ToolchainComponents {
            rustc,
            rustdoc,
            clippy,
            cargo,
            ..Default::default()
        };
        component.fill_libraries(host_libdir, target_libdir)?;
        Ok(component)
    }

    /// Finds known library components in the given `dir` and stores them in `self`.
    fn fill_libraries(&mut self, host_libdir: &Path, target_libdir: &Path) -> anyhow::Result<()> {
        let load_files = |path: &Path| -> anyhow::Result<Vec<(PathBuf, String)>> {
            let files = fs::read_dir(path)
                .with_context(|| {
                    format!(
                        "Cannot read lib dir `{}` to find components",
                        path.display()
                    )
                })?
                .map(|entry| Ok(entry?))
                .collect::<anyhow::Result<Vec<_>>>()?
                .into_iter()
                .filter(|entry| entry.path().is_file())
                .filter_map(|entry| {
                    entry
                        .path()
                        .file_name()
                        .and_then(|s| s.to_str())
                        .map(|s| (entry.path(), s.to_string()))
                })
                .collect();
            Ok(files)
        };

        // Look for librustc_driver.so and libLLVM.so in the *host* libdir
        let host_files = load_files(host_libdir)?;
        for (path, filename) in &host_files {
            if path.extension() == Some(OsStr::new("so")) && filename.starts_with("librustc_driver")
            {
                self.lib_rustc = Some(path.clone());
            }
        }

        // In older toolchains, the LLVM library is stored as libLLVM-<version>.so
        // In newer ones, this file is only a linker shim that actually redirects to
        // libLLVM.so.<version>.
        // So we need to check if we have the new name, and use it.
        // If not, we want to look up the original name.
        let new_llvm = host_files
            .iter()
            .find(|(_, filename)| filename.starts_with("libLLVM.so"));
        let old_llvm = host_files.iter().find(|(path, filename)| {
            path.extension() == Some(OsStr::new("so")) && filename.starts_with("libLLVM")
        });
        self.lib_llvm = new_llvm.or(old_llvm).map(|(path, _)| path.clone());

        // Now find libstd in the *target* libdir
        let target_files = load_files(target_libdir)?;
        for (path, filename) in target_files {
            if path.extension() == Some(OsStr::new("so")) && filename.starts_with("libstd") {
                self.lib_std = Some(path.clone());
            }
        }

        Ok(())
    }
}

#[derive(Clone, Copy, Default)]
pub struct ToolchainConfig<'a> {
    rustdoc: Option<&'a Path>,
    clippy: Option<&'a Path>,
    cargo: Option<&'a Path>,
    /// For `cargo --config <value>`.
    cargo_configs: &'a [String],
    id: Option<&'a str>,
}

impl<'a> ToolchainConfig<'a> {
    pub fn rustdoc(&mut self, rustdoc: Option<&'a Path>) -> &mut Self {
        self.rustdoc = rustdoc;
        self
    }

    pub fn clippy(&mut self, clippy: Option<&'a Path>) -> &mut Self {
        self.clippy = clippy;
        self
    }

    pub fn cargo(&mut self, cargo: Option<&'a Path>, configs: &'a [String]) -> &mut Self {
        self.cargo = cargo;
        self.cargo_configs = configs;
        self
    }

    pub fn id(&mut self, id: Option<&'a str>) -> &mut Self {
        self.id = id;
        self
    }
}

/// Get a toolchain from the input.
/// - `rustc`: check if the given one is acceptable.
/// - `rustdoc`: if one is given, check if it is acceptable. Otherwise, if
///   the `Doc` profile is requested, look for one next to the given `rustc`.
/// - `cargo`: if one is given, check if it is acceptable. Otherwise, look
///   for the nightly Cargo via `rustup`.
pub fn get_local_toolchain(
    profiles: &[Profile],
    codegen_backends: &[CodegenBackend],
    rustc: &str,
    toolchain_config: ToolchainConfig<'_>,
    id_suffix: &str,
    target_triple: String,
) -> anyhow::Result<Toolchain> {
    // `+`-prefixed rustc is an indicator to fetch the rustc of the toolchain
    // specified. This follows the similar pattern used by rustup's binaries
    // (e.g., `rustc +stage1`).
    let (rustc, id) = if let Some(toolchain) = rustc.strip_prefix('+') {
        let output = Command::new("rustup")
            .args(["which", "rustc", "--toolchain", toolchain])
            .output()
            .context("failed to run `rustup which rustc`")?;

        // Looks like a commit hash? Try to install it...
        if !output.status.success() && toolchain.len() == 40 {
            // No such toolchain exists, so let's try to install it with
            // rustup-toolchain-install-master.

            if Command::new("rustup-toolchain-install-master")
                .arg("-V")
                .output()
                .is_err()
            {
                anyhow::bail!("rustup-toolchain-install-master is not installed but must be");
            }

            let mut additional_components = vec![];
            if codegen_backends.contains(&CodegenBackend::Cranelift) {
                additional_components.push("rustc-codegen-cranelift");
            }

            let mut cmd = Command::new("rustup-toolchain-install-master");
            cmd.arg(toolchain);
            for component in additional_components {
                cmd.arg("-c").arg(component);
            }

            if !cmd
                .status()
                .context("failed to run `rustup-toolchain-install-master`")?
                .success()
            {
                anyhow::bail!(
                    "commit-like toolchain {} did not install successfully",
                    toolchain
                )
            }
        }

        let output = Command::new("rustup")
            .args(["which", "rustc", "--toolchain", toolchain])
            .output()
            .context("failed to run `rustup which rustc`")?;

        anyhow::ensure!(
            output.status.success(),
            "did not manage to obtain toolchain {toolchain}"
        );

        let s = String::from_utf8(output.stdout)
            .context("failed to convert `rustup which rustc` output to utf8")?;

        let rustc = PathBuf::from(s.trim());
        debug!("found rustc: {:?}", &rustc);

        // When the id comes from a +toolchain, the suffix is *not* added.
        let id = if let Some(id) = toolchain_config.id {
            let mut id = id.to_owned();
            id.push_str(id_suffix);
            id
        } else {
            toolchain.to_owned()
        };
        (rustc, id)
    } else {
        let rustc = PathBuf::from(rustc)
            .canonicalize()
            .with_context(|| format!("failed to canonicalize rustc executable {rustc:?}"))?;

        // When specifying rustc via a path, the suffix is always added to the
        // id.
        let mut id = if let Some(id) = toolchain_config.id {
            id.to_owned()
        } else {
            "Id".to_string()
        };
        id.push_str(id_suffix);

        (rustc, id)
    };

    let rustdoc =
        if let Some(rustdoc) = &toolchain_config.rustdoc {
            Some(rustdoc.canonicalize().with_context(|| {
                format!("failed to canonicalize rustdoc executable {rustdoc:?}")
            })?)
        } else if profiles.iter().any(|p| p.is_doc()) {
            // We need a `rustdoc`. Look for one next to `rustc`.
            if let Ok(rustdoc) = rustc.with_file_name("rustdoc").canonicalize() {
                debug!("found rustdoc: {:?}", &rustdoc);
                Some(rustdoc)
            } else {
                anyhow::bail!(
                    "'Doc' or 'DocJson' build specified but '--rustdoc' not specified and no \
                    'rustdoc' found next to 'rustc'"
                );
            }
        } else {
            // No `rustdoc` provided, but none needed.
            None
        };

    let clippy = if let Some(clippy) = &toolchain_config.clippy {
        Some(
            clippy
                .canonicalize()
                .with_context(|| format!("failed to canonicalize clippy executable {clippy:?}"))?,
        )
    } else if profiles.contains(&Profile::Clippy) {
        // We need a `clippy`. Look for one next to `rustc`.
        if let Ok(clippy) = rustc.with_file_name("clippy-driver").canonicalize() {
            debug!("found clippy: {:?}", &clippy);
            Some(clippy)
        } else {
            anyhow::bail!(
                    "'Clippy' build specified but '--clippy' not specified and no 'clippy-driver' found \
                    next to 'rustc'"
                );
        }
    } else {
        // No `clippy` provided, but none needed.
        None
    };
    let cargo = if let Some(cargo) = &toolchain_config.cargo {
        cargo
            .canonicalize()
            .with_context(|| format!("failed to canonicalize cargo executable {cargo:?}"))?
    } else {
        // Use the nightly cargo from `rustup`.
        let output = Command::new("rustup")
            .args(["which", "cargo", "--toolchain=nightly"])
            .output()
            .context("failed to run `rustup which cargo --toolchain=nightly`")?;
        if !output.status.success() {
            anyhow::bail!(
                "`rustup which cargo --toolchain=nightly` exited with status {}\nstderr={}",
                output.status,
                String::from_utf8_lossy(&output.stderr)
            )
        }
        let s = String::from_utf8(output.stdout)
            .context("failed to convert `rustup which cargo --toolchain=nightly` output to utf8")?;

        let cargo = PathBuf::from(s.trim());
        debug!("found cargo: {:?}", &cargo);
        cargo
    };
    let host_lib_dir = get_lib_dir_from_rustc(&rustc).context("Cannot find libdir for rustc")?;
    let target_lib_dir = target_libdir_from_host_libdir(&host_lib_dir, &target_triple);

    let mut components = ToolchainComponents::from_binaries_and_libdir(
        rustc,
        rustdoc,
        clippy,
        cargo,
        &host_lib_dir,
        &target_lib_dir,
    )?;
    components.cargo_configs = toolchain_config.cargo_configs.to_vec();
    Ok(Toolchain {
        components,
        id,
        triple: target_triple,
    })
}

/// Creates a toolchain from a *published* toolchain downloaded by rustup.
pub fn create_toolchain_from_published_version(
    toolchain: &str,
    target_triple: &str,
) -> anyhow::Result<Toolchain> {
    let status = Command::new("rustup")
        .args(["install", "--profile=minimal", toolchain])
        .status()
        .context("rustup install")?;
    if !status.success() {
        return Err(anyhow::anyhow!(
            "failed to install toolchain for {toolchain}",
        ));
    }

    let which = |tool| -> anyhow::Result<PathBuf> {
        let path = String::from_utf8(
            Command::new("rustup")
                .arg("which")
                .arg("--toolchain")
                .arg(toolchain)
                .arg(tool)
                .output()
                .context(format!("rustup which {tool}"))?
                .stdout,
        )
        .context("utf8")?;
        Ok(PathBuf::from(path.trim()))
    };
    let rustc = which("rustc")?;
    let rustdoc = which("rustdoc")?;
    let clippy = which("clippy")?;
    let cargo = which("cargo")?;

    debug!("Found rustc: {}", rustc.display());
    debug!("Found rustdoc: {}", rustdoc.display());
    debug!("Found clippy: {}", clippy.display());
    debug!("Found cargo: {}", cargo.display());

    let host_lib_dir = get_lib_dir_from_rustc(&rustc)?;
    let target_lib_dir = target_libdir_from_host_libdir(&host_lib_dir, target_triple);

    let components = ToolchainComponents::from_binaries_and_libdir(
        rustc,
        Some(rustdoc),
        Some(clippy),
        cargo,
        &host_lib_dir,
        &target_lib_dir,
    )?;

    Ok(Toolchain {
        components,
        id: toolchain.to_string(),
        triple: target_triple.to_string(),
    })
}

fn get_lib_dir_from_rustc(rustc: &Path) -> anyhow::Result<PathBuf> {
    let output = Command::new(rustc).arg("--print").arg("sysroot").output()?;
    if !output.status.success() {
        anyhow::bail!(
            "rustc failed to provide sysroot, exit status: {}\nstderr: {}",
            output.status,
            String::from_utf8_lossy(&output.stderr)
        );
    }
    let sysroot_path = String::from_utf8_lossy(&output.stdout);
    let lib_dir = Path::new(sysroot_path.as_ref().trim()).join("lib");
    if !lib_dir.exists() {
        anyhow::bail!(
            "rustc returned non-existent sysroot: `{}`",
            lib_dir.display()
        );
    }
    Ok(lib_dir)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs::File;

    #[test]
    fn fill_libraries() {
        let mut components = ToolchainComponents::default();

        let temp_dir: tempfile::TempDir = tempfile::tempdir().unwrap();
        let host_libdir = temp_dir.path();
        let target_libdir = target_libdir_from_host_libdir(host_libdir, "foo");
        std::fs::create_dir_all(&target_libdir).unwrap();

        let lib_rustc_path = create_lib(host_libdir, "librustc_driver.so");
        let lib_std_path = create_lib(
            &host_libdir.join("rustlib").join("foo").join("lib"),
            "libstd.so",
        );
        let lib_new_llvm_path = create_lib(host_libdir, "libLLVM.so.18.1-rust-1.78.0-nightly");

        components
            .fill_libraries(host_libdir, &target_libdir)
            .unwrap();

        assert_eq!(components.lib_rustc, Some(lib_rustc_path));
        assert_eq!(components.lib_std, Some(lib_std_path));
        assert_eq!(components.lib_llvm, Some(lib_new_llvm_path));
    }

    #[test]
    fn fill_old_llvm_library() {
        let mut components = ToolchainComponents::default();
        let lib_old_llvm = "libLLVM-17-rust-1.76.0-stable.so";

        // create mock dir and libraries
        let temp_dir: tempfile::TempDir = tempfile::tempdir().unwrap();
        let host_libdir = temp_dir.path();
        let target_libdir = target_libdir_from_host_libdir(host_libdir, "foo");
        std::fs::create_dir_all(&target_libdir).unwrap();

        let lib_old_llvm_path = create_lib(host_libdir, lib_old_llvm);

        components
            .fill_libraries(
                host_libdir,
                &target_libdir_from_host_libdir(temp_dir.path(), "foo"),
            )
            .unwrap();

        assert_eq!(components.lib_llvm, Some(lib_old_llvm_path));
    }

    fn create_lib(path: &Path, lib_name: &str) -> PathBuf {
        let lib_path = path.join(lib_name);
        File::create(&lib_path).unwrap();
        lib_path
    }
}
